%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                                                                         %
%              Joint Distribution of Productivity and Gaps                %
%                                                                         %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [FNC, F1, M, Lam, Deltagrid, qgrid_adj] = Joint_Dist(p)

% ---------------------------- Description --------------------------------

    % This function computes:

        % (i) The marginal distribution of productivity for varieties
        % without competitors. This distribution is referred to as
        % F^{NC}(q).

        % (ii) The joint distribution of productivity and gaps for varities
        % with competitors. This distribution is referred to as
        % F^{C}(q, Delta).

    % Finally, it uses these distributions to compute the misallocation
    % terms, M and Lambda.

    % For an overview of this function, see Section (RP-2.3.2).

% -------------------------------------------------------------------------
%                     1. Distribution without Competitors
% -------------------------------------------------------------------------

    % Constant that shows up in the differential equations - see Section (RP-2.3.2)
    d = (p.s-1)/(p.qbar-1);

    % Transformed productivity equivalent to minimum relative efficiency
    MinEff = log(p.wmin);

    % If the drift of the differential equation is positive...
    if p.qbar < 1

        % Determines the constant of integration - see Equation (RP-18)
        C = -p.varrho*(1-p.alpha)*p.wmin^(-d)/(p.varrho+d);

        % Solution to the differential equation - see Equation (RP-17)
        FNC_func = @(z) (C*exp(d*z) + (1-p.alpha) - d*(1-p.alpha)*p.wmin^(p.varrho)*exp(-p.varrho*z)/(p.varrho+d)) .* (z > MinEff);

    % If the drift of the differential equation is negative...
    else

        % Determines the constant of integration
        C = 0;

        % Solution to the differential equation - see Equation (RP-17)
        FNC_func = @(z) max((1-p.alpha) - d*(1-p.alpha)*p.wmin^(p.varrho)*exp(-p.varrho*z)/(p.varrho+d), 0);

    end

    % Evaluate the distribution over the exogenous grid
    FNC = FNC_func(p.qgrid);

% -------------------------------------------------------------------------
%                     2. Distribution with Competitors
% -------------------------------------------------------------------------

    % For the computation of the distribution for varieties with
    % competitors, we separate the problems into two stages:
        
        % (a) First, we obtain the marginal distribution of transformed
        % productivity.

        % (b) Second, we obtain the joint distribution of transformed
        % productivity and productivity gaps.

% ------------------- Marginal Distribution of Productivity ---------------

    % Compute the grid step implied in the productivity grid
    delta_q = p.qgrid(2) - p.qgrid(1);

    % Vector to store the solution for the marginal distribution
    Fmarginal = zeros(1, p.n3);

    % Vector that indicates whether q - log(lambda) is on the grid
    Lag_indicator = (p.qgrid - log(p.l) > p.qgrid(1));

    % variable that indicates the first "one" on the above vector
    Lag_index = find(Lag_indicator == 1, 1,'first');

    % ------------------- Case 1: Positive Drift --------------------------

    % If the drift of the differential equation is positive...
    if p.qbar < 1

        % Before Lag_index, Fmarginal is zero...
        for q = (Lag_index + 1):p.n3

        % If the argument of FNC is above MinEff, we compute normally
        if p.qgrid(q-1) - log(p.l) > MinEff

            % Term where the lags appear - see Equation (RP-19)
            Lag_term = FNC_func(p.qgrid(q-1) - log(p.l)) + interp1(p.qgrid(1:q-1), Fmarginal(1:q-1), p.qgrid(q-1) - log(p.l), 'linear');

        % If the argument of FNC is below MinEff, we forget (lagged) FNC
        else

            % Term where the lags appear - see Equation (RP-19)
            Lag_term = interp1(p.qgrid(1:q-1), Fmarginal(1:q-1), p.qgrid(q-1) - log(p.l), 'linear');

        end

        % Compute the derivative with respect to q - see Equation (RP-19)
        dfdq = d*Fmarginal(q-1) - p.alpha*d*Lag_term;

        % Compute Fmarginal evaluated at the next grid point
        Fmarginal(q) = Fmarginal(q-1) + dfdq*delta_q;

        end

    % ------------------- Case 2: Negative Drift --------------------------

    % The case of negative drift never showed up for our parameters. But if
    % that was the case, a different approach is needed to solve the
    % differential equation. This method would be significantly slower than
    % the one employed for the case of positive drifts.

    % ---------------------------------------------------------------------

    % If the drift of the differential equation is negative...     
    else

        % Determine an initial guess for the marginal distribution
        Fmarginal_old  = FNC*p.alpha/(1-p.alpha);

        % Preallocate a vector to store the updated distribution
        Fmarginal_new = Fmarginal_old;

        % Compute the lagged version of the NC distribution
        Lag_NC = FNC_func(p.qgrid - log(p.l));

        % Start an error tracker at a large enough value
        err = 1;

        % While the error is greater than some tolerance
        while err > 10e-6

            % Compute the lagged version of the marginal distribution
            Lag_C = @(x) interp1(p.qgrid, Fmarginal_old, x, "spline", 0);

            % Loop over possible productivity values...
            for q = fliplr(2:p.n3)

                % Evaluate the derivative with respect to q
                dfdq = d*Fmarginal_new(q) - d*p.alpha* (Lag_C(p.qgrid(q)-log(p.l)) + Lag_NC(q));

                % Update the new version of the distribution
                Fmarginal_new(q-1) = Fmarginal_new(q) - dfdq * delta_q;

            end

            % Compute the error at the current iteration
            err = max(abs(Fmarginal_new-Fmarginal_old), [], "all");

            % Update the old distribution with the updated one
            Fmarginal_old = Fmarginal_new;

        end

        % Assign the marginal distribution to the appropriate vector
        Fmarginal = Fmarginal_old;

    end

% --------------- Joint Distribution of Gaps and Productivity -------------

    % Determines the upper bound on the grid for productivity gaps
    Dlim = 4;

    % Determines the number of points on the grid for productivity gaps
    Deltagrid_pts = 250; 

    % Creates a (non-linear) grid for productivity gaps starting on lambda
    Deltagrid = exp(linspace(log(p.l), log(Dlim), Deltagrid_pts));
    
    % Preallocate a matrix to store the joint distribution
    F1 = zeros(Deltagrid_pts, p.n3);

    % Attributes the marginal distribution to the highest value of Delta
    F1(end,:) = Fmarginal;

    % We interpolate the marginal distribution to have lags for every q
    Fmarginal_lags_short = interp1(p.qgrid, Fmarginal, p.qgrid(Lag_index:end) - log(p.l), 'spline');

    % Complete the vector w/ zeros for arguments that fall before the grid
    Fmarginal_lags = [zeros(1, Lag_index - 1), Fmarginal_lags_short];

    % We compute the lagged version of the FNC distribution
    FNC_lags = FNC_func(p.qgrid - log(p.l));

    % NaNs where p.qgrid - log(p.l) is below MinEff should be zero
    FNC_lags(isnan(FNC_lags)) = 0;

    % For each value of productivity gap...
    for i = 2:Deltagrid_pts

        % If the drift of the differential equation is positive...
        if p.qbar < 1

            % ... we compute the backward derivative wrt q
            dfdq = [0, diff(F1(i-1, :))] / delta_q;

        % If the drift of the differential equation is positive...
        else

            % ... we compute the forward derivative wrt q
            dfdq = [diff(F1(i-1, :)), 0] / delta_q;

        end

        % Compute the derivative of F wrt Delta - see Equation (RP-20)
        dfdd = (p.eta+p.d0)/(1-p.alpha)*(dfdq/d - F1(i-1, :) + p.alpha*(Fmarginal_lags+FNC_lags))/(p.In*Deltagrid(i-1));

        % Compute joint distribution evaluated at the next gap grid point
        F1(i, :)= F1(i-1, :) + dfdd*(Deltagrid(i)-Deltagrid(i-1));
        
    end

% -------------------------------------------------------------------------
%                       3. Recovering the Densities
% -------------------------------------------------------------------------

    % Given the distributions derived in the Sections above, the densities
    % would be defined over log(q/Q). In this Section, we transform into
    % densities defined over q, noting that Q can be normalized

    % ---------------------------------------------------------------------

    % Determines the grid for q that is implied in the grid for log(q)
    Q_endog_grid = exp(p.qgrid);

    % Obtains density for varieties without competitors defined over log(q)
    FNCdens = zeros(1, p.n3);
    FNCdens(1:end-1) = diff(FNC)/delta_q;

    % Transforms this density to be defined over q
    FNCdens_endog = FNCdens ./ Q_endog_grid;
    FNCdens_endog(isnan(FNCdens_endog)) = 0;

    % Obtains density for varieties with competitors defined over log(q)
    FdensC = zeros(Deltagrid_pts, p.n3);
    FdensC(:, 1:end-1) = diff(F1, 1, 2)/delta_q;
    
    % Transforms this density to be defined over q
    FCdens_endog = FdensC ./ repmat(Q_endog_grid, Deltagrid_pts, 1);

    % Also, for every value of productivity...
    for q = 1:p.n3

        % ...we need to take differences over the gap dimension
        FCdens_endog(:, q) = gradient(FCdens_endog(:, q), Deltagrid);

    end

% -------------------------------------------------------------------------
%                          4. Numerical Issues
% -------------------------------------------------------------------------

    % Compute exponential terms that enter the integrals - see Equation (RP-21)
    exp_term = repmat(Q_endog_grid.^(p.s-1), Deltagrid_pts, 1);

    % Compute the adjustment factor for the integrals  - see Equation (RP-21)
    Adj_NC = trapz(Q_endog_grid, FNCdens_endog .* Q_endog_grid.^(p.s-1));
    Adj_C = trapz(Q_endog_grid, trapz(Deltagrid, FCdens_endog.*exp_term, 1));
    Adj_fac = Adj_NC + Adj_C;

    % Constant that shifts the grid by making Equation (RP-21) hold
    Adj_constant = -log(Adj_fac)/(p.s-1);

    % Shifts the productivity grid accordingly
    qgrid_adj = p.qgrid + Adj_constant;

    % ---------------------- Repeats Section 3 ----------------------------

    % Here, we just recover the densities once again. But this time, the
    % relevant grid is the one already adjusted by the previous step.

    % ---------------------------------------------------------------------

    % Determines the grid for q that is implied in the grid for log(q)
    Q_endog_grid = exp(qgrid_adj);

    % Obtains density for varieties without competitors defined over log(q)
    FNCdens = zeros(1, p.n3);
    FNCdens(1:end-1) = diff(FNC)/delta_q;

    % Transforms this density to be defined over q
    FNCdens_endog = FNCdens ./ Q_endog_grid;
    FNCdens_endog(isnan(FNCdens_endog)) = 0;

    % Obtains density for varieties with competitors defined over log(q)
    FdensC = zeros(Deltagrid_pts, p.n3);
    FdensC(:, 1:end-1) = diff(F1, 1, 2)/delta_q;

    % Transforms this density to be defined over q
    FCdens_endog = FdensC ./ repmat(Q_endog_grid, Deltagrid_pts, 1);

    % Also, for every value of productivity...
    for q = 1:p.n3

        % ...we need to take differences over the gap dimension
        FCdens_endog(:, q) = gradient(FCdens_endog(:, q), Deltagrid);

    end

% -------------------------------------------------------------------------
%                          5. Misallocation Terms
% -------------------------------------------------------------------------

    % Compute the mark-up and exponential terms that enter the integrals - see Equations (RP-13) and (RP-14)
    MarkupTerm = repmat(min(p.mu, Deltagrid),p.n3,1)';
    exp_term = repmat(Q_endog_grid.^(p.s-1), Deltagrid_pts, 1);

    % Compute the integral term associated with no competitors - see Equations (RP-13) and (RP-14)
    Int_NC = trapz(Q_endog_grid, FNCdens_endog .* Q_endog_grid.^(p.s-1));

    % Compute the first integrals associated varieties with competitors - see Equations (RP-13) and (RP-14)
    integrand = FCdens_endog.*exp_term.*MarkupTerm.^(1-p.s);
    Int_C1 = trapz(Q_endog_grid, trapz(Deltagrid, integrand, 1));

    % Compute the second integrals associated varieties with competitors - see Equations (RP-13) and (RP-14)
    integrand = FCdens_endog.*exp_term.*MarkupTerm.^(-p.s);
    Int_C2 = trapz(Q_endog_grid, trapz(Deltagrid, integrand, 1));

    % Compute the misallocation terms - see Equations (RP-13) and (RP-14)
    Lam = (Int_C2 + p.mu^(-p.s)*Int_NC)/(Int_C1 + p.mu^(1-p.s)*Int_NC);
    M = (Int_C1 + p.mu^(1-p.s)*Int_NC)^(1/(p.s-1))/Lam;

end



